import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from enum import Enum

# 飞行模式枚举
class FlightMode(Enum):
    TAKEOFF = 0
    WAYPOINT = 1
    RETURN_HOME = 2
    LANDING = 3

# 测试模式枚举
class TestMode(Enum):
    TAKEOFF_TEST = 0
    LANDING_TEST = 1
    WAYPOINT_TEST = 2
    RETURN_HOME_TEST = 3
    COMPLETE_TEST = 4

# 四旋翼参数
MASS = 1.0         # kg
GRAVITY = 9.81     # m/s^2
DT = 0.01          # 时间步长
SIM_TIME = 30      # 仿真时长（秒）
INERTIA = np.array([0.1, 0.1, 0.1])  # 转动惯量

# 飞行参数
TAKEOFF_HEIGHT = 20.0  # 起飞高度
LANDING_SPEED = 0.5   # 降落速度
WAYPOINT_RADIUS = 0.1 # 到达航路点的判定半径
HOME_POSITION = np.array([0.0, 0.0, 0.0])  # 返航点位置

# 航路点列表
WAYPOINTS = [
     np.array([0.0, 0.0, 20.0]),
     np.array([20.0, 20.0, 20.0]),  # 右前上方
    np.array([20.0, -20.0, 20.0]), # 右后上方
    np.array([-20.0, -20.0, 20.0]),# 左后上方
    np.array([-20.0, 20.0, 20.0]),  # 左前上方
    np.array([-120.0, 20.0, 20.0])
]

# 四旋翼几何参数
ARM_LENGTH = 0.1   # 旋翼臂长
ROTOR_RADIUS = 0.05  # 旋翼半径
AXIS_LENGTH = 0.3  # 姿态指示器长度

# 测试参数
TEST_COMPLETION_TIME = 5.0  # 每个测试模式完成后等待时间
test_mode = TestMode.WAYPOINT_TEST  # 当前测试模式
test_completion_time = 0  # 测试完成计时器

# 速度控制参数
MAX_SPEED = 15.0    # 最大飞行速度 (m/s)
APPROACH_DISTANCE = 2.0  # 接近航路点时开始减速的距离



def create_attitude_indicators(pos, att):
    """创建姿态指示器"""
    # 创建基础坐标轴
    axes = np.array([
        [AXIS_LENGTH, 0, 0],    # X轴
        [0, AXIS_LENGTH, 0],    # Y轴
        [0, 0, AXIS_LENGTH],    # Z轴
    ])
    
    # 旋转矩阵
    roll, pitch, yaw = att
    Rx = np.array([[1, 0, 0],
                   [0, np.cos(roll), -np.sin(roll)],
                   [0, np.sin(roll), np.cos(roll)]])
    
    Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)],
                   [0, 1, 0],
                   [-np.sin(pitch), 0, np.cos(pitch)]])
    
    Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0],
                   [np.sin(yaw), np.cos(yaw), 0],
                   [0, 0, 1]])
    
    R = Rz @ Ry @ Rx
    
    # 应用旋转
    axes = axes @ R.T
    
    # 添加位置偏移
    axes += pos
    
    return axes

def create_quadcopter_faces(vertices):
    """创建四旋翼的面"""
    # 创建中心点
    center = np.mean(vertices, axis=0)
    
    # 创建连接中心点和各个顶点的三角形
    faces = []
    for i in range(4):
        faces.append([center, vertices[i], vertices[(i+1)%4]])
    
    return faces

# PID 控制器
class PID:
    def __init__(self, kp, ki, kd, setpoint=0):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.setpoint = setpoint
        self.integral = 0
        self.last_error = 0

    def update(self, value, dt):
        error = self.setpoint - value
        self.integral += error * dt
        derivative = (error - self.last_error) / dt
        self.last_error = error
        return self.kp * error + self.ki * self.integral + self.kd * derivative

    def set_setpoint(self, setpoint):
        self.setpoint = setpoint

# 状态变量
pos = np.array([0.0, 0.0, 0.0])   # 位置 x, y, z
vel = np.array([0.0, 0.0, 0.0])   # 速度 vx, vy, vz
att = np.array([0.0, 0.0, 0.0])   # 姿态角 roll, pitch, yaw
ang_vel = np.array([0.0, 0.0, 0.0])  # 角速度

# 飞行模式
current_mode = FlightMode.TAKEOFF
current_waypoint_index = 0  # 提前声明

# 位置PID控制器
pid_x = PID(2.0, 0.0, 1.0)
pid_y = PID(2.0, 0.0, 1.0)
pid_z = PID(10.0, 2.0, 5.0)

# 速度PID控制器
pid_vx = PID(2.0, 0.0, 0.5)
pid_vy = PID(2.0, 0.0, 0.5)
pid_vz = PID(2.0, 0.0, 0.5)

# 姿态控制器
pid_roll = PID(2.0, 0.0, 0.5)
pid_pitch = PID(2.0, 0.0, 0.5)
pid_yaw = PID(2.0, 0.0, 0.5)

# 确保起飞测试时有目标高度
if test_mode == TestMode.TAKEOFF_TEST:
    pid_z.set_setpoint(TAKEOFF_HEIGHT)
# 确保航路点测试时有目标航点
if test_mode == TestMode.WAYPOINT_TEST:
    current_mode = FlightMode.WAYPOINT
    current_waypoint_index = 0
    pid_x.set_setpoint(WAYPOINTS[0][0])
    pid_y.set_setpoint(WAYPOINTS[0][1])
    pid_z.set_setpoint(WAYPOINTS[0][2])

# 存储历史数据
positions = []
attitudes = []
flight_modes = []
velocities = []

def calculate_velocity_setpoint(target_pos, current_pos):
    """计算速度设定点，考虑最大速度和接近目标时的减速"""
    # 计算到目标的向量
    direction = target_pos - current_pos
    distance = np.linalg.norm(direction)
    
    if distance < 0.01:  # 防止除以零
        return np.zeros(3)
    
    # 单位方向向量
    direction_unit = direction / distance
    
    # 根据距离计算目标速度
    if distance > APPROACH_DISTANCE:
        # 远距离时使用最大速度
        return direction_unit * MAX_SPEED
    else:
        # 接近目标时线性减速
        speed_factor = distance / APPROACH_DISTANCE
        return direction_unit * MAX_SPEED * speed_factor

def update_test_mode():
    global test_mode, current_mode, test_completion_time, pos, vel, att, ang_vel, current_waypoint_index
    
    if test_completion_time > 0:
        test_completion_time -= DT
        return
    
    if test_mode == TestMode.TAKEOFF_TEST:
        if current_mode == FlightMode.TAKEOFF and pos[2] >= TAKEOFF_HEIGHT:
            test_mode = TestMode.LANDING_TEST
            current_mode = FlightMode.LANDING
            pid_z.set_setpoint(0)
            test_completion_time = TEST_COMPLETION_TIME
    
    elif test_mode == TestMode.LANDING_TEST:
        if current_mode == FlightMode.LANDING and pos[2] <= 0.1:
            test_mode = TestMode.WAYPOINT_TEST
            current_mode = FlightMode.TAKEOFF
            current_waypoint_index = 0
            pid_x.set_setpoint(WAYPOINTS[0][0])
            pid_y.set_setpoint(WAYPOINTS[0][1])
            pid_z.set_setpoint(WAYPOINTS[0][2])
            test_completion_time = TEST_COMPLETION_TIME
    
    elif test_mode == TestMode.WAYPOINT_TEST:
        if current_mode == FlightMode.WAYPOINT and current_waypoint_index >= len(WAYPOINTS):
            test_mode = TestMode.RETURN_HOME_TEST
            current_mode = FlightMode.RETURN_HOME
            pid_x.set_setpoint(HOME_POSITION[0])
            pid_y.set_setpoint(HOME_POSITION[1])
            pid_z.set_setpoint(HOME_POSITION[2])
            test_completion_time = TEST_COMPLETION_TIME
    
    elif test_mode == TestMode.RETURN_HOME_TEST:
        if current_mode == FlightMode.RETURN_HOME and np.linalg.norm(pos - HOME_POSITION) < WAYPOINT_RADIUS:
            test_mode = TestMode.COMPLETE_TEST
            current_mode = FlightMode.LANDING
            pid_z.set_setpoint(0)
            test_completion_time = TEST_COMPLETION_TIME

def update_flight_mode():
    global current_mode, current_waypoint_index
    
    if current_mode == FlightMode.TAKEOFF:
        if pos[2] >= TAKEOFF_HEIGHT:
            if test_mode == TestMode.TAKEOFF_TEST:
                return
            current_mode = FlightMode.WAYPOINT
            current_waypoint_index = 0
            pid_x.set_setpoint(WAYPOINTS[0][0])
            pid_y.set_setpoint(WAYPOINTS[0][1])
            pid_z.set_setpoint(WAYPOINTS[0][2])
    
    elif current_mode == FlightMode.WAYPOINT:
        current_waypoint = WAYPOINTS[current_waypoint_index]
        distance = np.linalg.norm(pos - current_waypoint)
        
        if distance < WAYPOINT_RADIUS:
            current_waypoint_index += 1
            if current_waypoint_index >= len(WAYPOINTS):
                if test_mode == TestMode.WAYPOINT_TEST:
                    return
                current_mode = FlightMode.RETURN_HOME
                pid_x.set_setpoint(HOME_POSITION[0])
                pid_y.set_setpoint(HOME_POSITION[1])
                pid_z.set_setpoint(HOME_POSITION[2])
            else:
                next_waypoint = WAYPOINTS[current_waypoint_index]
                pid_x.set_setpoint(next_waypoint[0])
                pid_y.set_setpoint(next_waypoint[1])
                pid_z.set_setpoint(next_waypoint[2])
    
    elif current_mode == FlightMode.RETURN_HOME:
        distance = np.linalg.norm(pos - HOME_POSITION)
        if distance < WAYPOINT_RADIUS:
            if test_mode == TestMode.RETURN_HOME_TEST:
                return
            current_mode = FlightMode.LANDING
            pid_z.set_setpoint(0)
    
    elif current_mode == FlightMode.LANDING:
        # 让 setpoint 以最大下降速度递减
        new_setpoint = max(pid_z.setpoint - LANDING_SPEED * DT, 0)
        pid_z.set_setpoint(new_setpoint)
        if pos[2] <= 0.1:
            if test_mode == TestMode.COMPLETE_TEST:
                return
            current_mode = FlightMode.TAKEOFF
#----------------------------------------------------------------------------
# 新增变量: 计算姿态信息
calculated_att = np.array([0.0, 0.0, 0.0])  # 计算得到的姿态角
prev_vel = np.array([0.0, 0.0, 0.0])       # 上一时刻速度
acc = np.array([0.0, 0.0, 0.0])            # 当前加速度
MAX_PITCH_ANGLE =35# np.radians(60)           # 最大俯仰角度 (30度)
# def calculate_attitude_from_motion(vel, acc, dt):
#     """根据速度和加速度计算姿态信息"""
#     global calculated_att
    
#     # 1. 计算航向角 (yaw) - 基于飞行方向
#     if np.linalg.norm(vel[:2]) > 0.1:  # 只有水平速度足够大时才计算
#         yaw =0.0 #np.arctan2(vel[1], vel[0])
#     else:
#         yaw =0.0 #calculated_att[2]  # 保持当前航向
    
#     # 2. 计算俯仰角 (pitch) - 基于加速度
#     if np.linalg.norm(acc) > 0.1:  # 只有加速度足够大时才计算
#         # 计算前向加速度分量 (在无人机坐标系中)
#         forward_acc = acc[0] * np.cos(yaw) + acc[1] * np.sin(yaw)
        
#         # 根据前向加速度计算俯仰角 (限制在±30度之间)
#         pitch = -np.clip(forward_acc / MAX_SPEED, -1.0, 1.0) * MAX_PITCH_ANGLE
#     else:
#         pitch = 0.0  # 无加速度时保持水平
    
#     # 3. 滚转角 (roll) - 暂时保持为0，可以根据需要添加转弯时的滚转
#     roll = 0.0
    
#     # 更新计算姿态
#     calculated_att = np.array([roll, pitch, yaw])
    
#     return calculated_att
MAX_PITCH_ANGLE = 20  # 最大俯仰角(度)
MAX_ROLL_ANGLE = 30   # 最大滚转角(度)
def calculate_attitude_from_motion(vel, acc, dt):
    """根据速度和加速度计算平滑的姿态信息"""
    global calculated_att, last_pitch_sign, last_yaw, filtered_pitch, filtered_yaw
    
    # 初始化全局变量
    if 'calculated_att' not in globals():
        calculated_att = np.array([0.0, 0.0, 0.0])
    if 'filtered_pitch' not in globals():
        filtered_pitch = 0.0
    if 'filtered_yaw' not in globals():
        filtered_yaw = calculated_att[2]
    
    # 滤波器参数
    alpha = 0.2  # 低通滤波系数 (0.1-0.3)
    max_pitch_rate = 10.0 * dt  # 最大俯仰角变化率 (度/秒)
    max_yaw_rate = 30.0 * dt    # 最大偏航角变化率 (度/秒)
    
    # 1. 计算航向角 (yaw) - 基于飞行方向
    if np.linalg.norm(vel[:2]) > 0.1:  # 只有水平速度足够大时才计算
        new_yaw = np.arctan2(vel[1], vel[0]) * 57.3  # 转换为度
        
        # 处理角度环绕问题 (从-180到180的跳变)
        if abs(new_yaw - filtered_yaw) > 180:
            if new_yaw > filtered_yaw:
                new_yaw -= 360
            else:
                new_yaw += 360
                
        # 应用变化率限制和低通滤波
        yaw_diff = np.clip(new_yaw - filtered_yaw, -max_yaw_rate, max_yaw_rate)
        filtered_yaw += yaw_diff
        filtered_yaw = alpha * new_yaw + (1 - alpha) * filtered_yaw
    else:
        # 保持当前航向
        filtered_yaw = calculated_att[2]
    
    # 2. 计算俯仰角 (pitch) - 基于加速度方向和速度大小
    # 计算前向加速度分量 (在无人机坐标系中)
    yaw_rad = np.radians(filtered_yaw)
    forward_accel = acc[0] * np.cos(yaw_rad) + acc[1] * np.sin(yaw_rad)
    
    # 计算前向速度分量 (在无人机坐标系中)
    forward_speed = vel[0] * np.cos(yaw_rad) + vel[1] * np.sin(yaw_rad)
    
    # 决定俯仰方向 - 基于加速度的方向
    pitch_direction = np.sign(forward_accel) if abs(forward_accel) > 0.1 else 0
    
    # 计算俯仰程度 - 基于速度的大小
    speed_factor = np.clip(abs(forward_speed) / MAX_SPEED, 0.0, 1.0)
    
    # 目标俯仰角 = 方向 * 程度 * 最大角度
    target_pitch = pitch_direction * speed_factor * MAX_PITCH_ANGLE
    
    # 应用变化率限制和低通滤波
    pitch_diff = np.clip(target_pitch - filtered_pitch, -max_pitch_rate, max_pitch_rate)
    filtered_pitch += pitch_diff
    filtered_pitch = alpha * target_pitch + (1 - alpha) * filtered_pitch
    
    # 3. 计算滚转角 (roll) - 基于转弯时的向心力
    if np.linalg.norm(vel[:2]) > 0.5:  # 只有速度足够大时才计算滚转
        # 计算转弯半径和向心加速度
        lateral_accel = -acc[0] * np.sin(yaw_rad) + acc[1] * np.cos(yaw_rad)
        target_roll = np.clip(np.arctan(lateral_accel / 9.81) * 57.3, -MAX_ROLL_ANGLE, MAX_ROLL_ANGLE)
    else:
        target_roll = 0.0
    
    # 滚转角也需要平滑
    roll_diff = np.clip(target_roll - calculated_att[0], -max_pitch_rate, max_pitch_rate)
    filtered_roll = calculated_att[0] + roll_diff
    
    # 更新计算姿态 (roll, pitch, yaw)
    calculated_att = np.array([
        alpha * filtered_roll + (1 - alpha) * calculated_att[0],
        filtered_pitch,
        filtered_yaw
    ])
    
    return calculated_att

#----------------------------------------------------------------------------
def create_quadcopter_model(pos, att, rotor_angle):
    """创建更精细的四旋翼模型"""
    # 创建旋转矩阵
    roll, pitch, yaw = att
    Rx = np.array([[1, 0, 0],
                   [0, np.cos(roll), -np.sin(roll)],
                   [0, np.sin(roll), np.cos(roll)]])
    
    Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)],
                   [0, 1, 0],
                   [-np.sin(pitch), 0, np.cos(pitch)]])
    
    Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0],
                   [np.sin(yaw), np.cos(yaw), 0],
                   [0, 0, 1]])
    
    R = Rz @ Ry @ Rx
    
    # 1. 机身（立方体）
    body_size = 0.1
    body_vertices = np.array([
        [-body_size/2, -body_size/2, -body_size/2],
        [body_size/2, -body_size/2, -body_size/2],
        [body_size/2, body_size/2, -body_size/2],
        [-body_size/2, body_size/2, -body_size/2],
        [-body_size/2, -body_size/2, body_size/2],
        [body_size/2, -body_size/2, body_size/2],
        [body_size/2, body_size/2, body_size/2],
        [-body_size/2, body_size/2, body_size/2]
    ])
    
    # 应用旋转
    body_vertices = body_vertices @ R.T
    # 添加位置偏移
    body_vertices += pos
    
    # 机身面
    body_faces = [
        [body_vertices[0], body_vertices[1], body_vertices[2], body_vertices[3]],  # 底面
        [body_vertices[4], body_vertices[5], body_vertices[6], body_vertices[7]],  # 顶面
        [body_vertices[0], body_vertices[1], body_vertices[5], body_vertices[4]],  # 前面
        [body_vertices[2], body_vertices[3], body_vertices[7], body_vertices[6]],  # 后面
        [body_vertices[0], body_vertices[3], body_vertices[7], body_vertices[4]],  # 左面
        [body_vertices[1], body_vertices[2], body_vertices[6], body_vertices[5]]   # 右面
    ]
    
    # 2. 旋翼臂
    arms = []
    arm_width = 0.02
    arm_height = 0.01
    
    arm_positions = [
        [ARM_LENGTH, 0, 0],  # 右前
        [0, ARM_LENGTH, 0],  # 前
        [-ARM_LENGTH, 0, 0], # 左前
        [0, -ARM_LENGTH, 0]  # 后
    ]
    
    for arm_pos in arm_positions:
        # 臂的顶点
        arm_vertices = np.array([
            [-arm_width/2, -arm_height/2, 0],
            [arm_width/2, -arm_height/2, 0],
            [arm_width/2, arm_height/2, 0],
            [-arm_width/2, arm_height/2, 0],
            [-arm_width/2, -arm_height/2, arm_pos[0]],
            [arm_width/2, -arm_height/2, arm_pos[0]],
            [arm_width/2, arm_height/2, arm_pos[0]],
            [-arm_width/2, arm_height/2, arm_pos[0]]
        ])
        
        # 旋转并移动到正确位置
        arm_vertices = arm_vertices @ R.T
        arm_vertices += pos
        
        # 臂的面
        arm_faces = [
            [arm_vertices[0], arm_vertices[1], arm_vertices[2], arm_vertices[3]],  # 底面
            [arm_vertices[4], arm_vertices[5], arm_vertices[6], arm_vertices[7]],  # 顶面
            [arm_vertices[0], arm_vertices[1], arm_vertices[5], arm_vertices[4]],  # 前面
            [arm_vertices[2], arm_vertices[3], arm_vertices[7], arm_vertices[6]],  # 后面
            [arm_vertices[0], arm_vertices[3], arm_vertices[7], arm_vertices[4]],  # 左面
            [arm_vertices[1], arm_vertices[2], arm_vertices[6], arm_vertices[5]]   # 右面
        ]
        
        arms.append(arm_faces)
    
    # 3. 旋翼
    rotors = []
    rotor_width = 0.1
    rotor_thickness = 0.01
    
    for i, arm_pos in enumerate(arm_positions):
        # 计算旋翼位置
        rotor_pos = np.array([
            pos[0] + arm_pos[0] * R[0,0] + arm_pos[1] * R[0,1] + arm_pos[2] * R[0,2],
            pos[1] + arm_pos[0] * R[1,0] + arm_pos[1] * R[1,1] + arm_pos[2] * R[1,2],
            pos[2] + arm_pos[0] * R[2,0] + arm_pos[1] * R[2,1] + arm_pos[2] * R[2,2]
        ])
        
        # 为每个旋翼创建不同的旋转速度（交替旋转方向）
        blade_angle = rotor_angle + i * np.pi/2
        
        # 创建两个旋翼叶片
        for blade in range(2):
            # 计算叶片的方向（互相垂直）
            blade_rot = blade_angle + blade * np.pi
            
            # 叶片顶点
            blade_vertices = np.array([
                [-rotor_thickness/2, -rotor_width/2, 0],
                [rotor_thickness/2, -rotor_width/2, 0],
                [rotor_thickness/2, rotor_width/2, 0],
                [-rotor_thickness/2, rotor_width/2, 0]
            ])
            
            # 旋转叶片
            R_blade = np.array([
                [np.cos(blade_rot), -np.sin(blade_rot), 0],
                [np.sin(blade_rot), np.cos(blade_rot), 0],
                [0, 0, 1]
            ])
            
            # 应用旋翼臂的旋转
            blade_vertices = blade_vertices @ R_blade.T @ R.T
            blade_vertices += rotor_pos
            
            # 叶片面
            blade_faces = [
                [blade_vertices[0], blade_vertices[1], blade_vertices[2], blade_vertices[3]]
            ]
            
            rotors.append(blade_faces)
    
    # 4. 起落架
    landing_gear = []
    gear_height = 0.15
    gear_width = 0.02
    
    for i, arm_pos in enumerate(arm_positions):
        # 计算起落架位置
        gear_pos = np.array([
            pos[0] + arm_pos[0] * R[0,0] + arm_pos[1] * R[0,1] + arm_pos[2] * R[0,2],
            pos[1] + arm_pos[0] * R[1,0] + arm_pos[1] * R[1,1] + arm_pos[2] * R[1,2],
            pos[2] + arm_pos[0] * R[2,0] + arm_pos[1] * R[2,1] + arm_pos[2] * R[2,2]
        ])
        
        # 起落架顶点
        gear_vertices = np.array([
            [-gear_width/2, -gear_width/2, 0],
            [gear_width/2, -gear_width/2, 0],
            [gear_width/2, gear_width/2, 0],
            [-gear_width/2, gear_width/2, 0],
            [-gear_width/2, -gear_width/2, -gear_height],
            [gear_width/2, -gear_width/2, -gear_height],
            [gear_width/2, gear_width/2, -gear_height],
            [-gear_width/2, gear_width/2, -gear_height]
        ])
        
        # 应用旋转
        gear_vertices = gear_vertices @ R.T
        gear_vertices += gear_pos
        
        # 起落架面
        gear_faces = [
            [gear_vertices[0], gear_vertices[1], gear_vertices[5], gear_vertices[4]],  # 前面
            [gear_vertices[2], gear_vertices[3], gear_vertices[7], gear_vertices[6]],  # 后面
            [gear_vertices[0], gear_vertices[3], gear_vertices[7], gear_vertices[4]],  # 左面
            [gear_vertices[1], gear_vertices[2], gear_vertices[6], gear_vertices[5]]   # 右面
        ]
        
        landing_gear.append(gear_faces)
    
    return body_faces, arms, rotors, landing_gear

# 创建图形
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(111, projection='3d')

# 设置坐标轴范围
ax.set_xlim([-100, 100])
ax.set_ylim([-100, 100])
ax.set_zlim([-100, 100])
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

# 绘制航路点
for waypoint in WAYPOINTS:
    ax.scatter(waypoint[0], waypoint[1], waypoint[2], color='green', marker='^')

# 绘制返航点
ax.scatter(HOME_POSITION[0], HOME_POSITION[1], HOME_POSITION[2], color='red', label='Home')

# 创建更精细的四旋翼模型
rotor_angle = 0
body_faces, arms, rotors, landing_gear = create_quadcopter_model(pos, att, rotor_angle)

# 创建四旋翼模型的多边形集合
body_collection = Poly3DCollection(body_faces, alpha=0.7, facecolor='blue', edgecolor='black')
ax.add_collection3d(body_collection)

arm_collections = []
for arm_faces in arms:
    arm_collection = Poly3DCollection(arm_faces, alpha=0.7, facecolor='darkgray', edgecolor='black')
    ax.add_collection3d(arm_collection)
    arm_collections.append(arm_collection)

rotor_collections = []
for rotor_faces in rotors:
    rotor_collection = Poly3DCollection(rotor_faces, alpha=0.5, facecolor='red', edgecolor='darkred')
    ax.add_collection3d(rotor_collection)
    rotor_collections.append(rotor_collection)

landing_gear_collections = []
for gear_faces in landing_gear:
    gear_collection = Poly3DCollection(gear_faces, alpha=0.7, facecolor='silver', edgecolor='black')
    ax.add_collection3d(gear_collection)
    landing_gear_collections.append(gear_collection)

# 创建姿态指示器
axes = create_attitude_indicators(pos, att)
x_axis, = ax.plot([pos[0], axes[0,0]], [pos[1], axes[0,1]], [pos[2], axes[0,2]], 'r-', linewidth=2)
y_axis, = ax.plot([pos[0], axes[1,0]], [pos[1], axes[1,1]], [pos[2], axes[1,2]], 'g-', linewidth=2)
z_axis, = ax.plot([pos[0], axes[2,0]], [pos[1], axes[2,1]], [pos[2], axes[2,2]], 'b-', linewidth=2)

# 创建轨迹线
trajectory_line, = ax.plot([], [], [], 'b-', alpha=0.5)

# 创建速度向量
velocity_vector, = ax.plot([pos[0], pos[0]+vel[0]*0.2], [pos[1], pos[1]+vel[1]*0.2], [pos[2], pos[2]+vel[2]*0.2], 'y-', linewidth=2)

# 添加模式文本
mode_text = ax.text2D(0.02, 0.95, '', transform=ax.transAxes)
speed_text = ax.text2D(0.02, 0.90, '', transform=ax.transAxes)

def update(frame):
    global pos, vel, att, ang_vel, rotor_angle

    # 随机风的影响
    wind = np.random.normal(0, 0.0001, size=3)  # 轻微的随机风
    vel += wind
    
    # 更新测试模式
    update_test_mode()
    
    # 更新飞行模式
    update_flight_mode()

    # 计算速度设定点
    if current_mode == FlightMode.TAKEOFF:
        vx_setpoint = 0
        vy_setpoint = 0
        vz_setpoint = min(MAX_SPEED, (TAKEOFF_HEIGHT - pos[2]) / 2)
    elif current_mode == FlightMode.WAYPOINT:
        target_pos = WAYPOINTS[current_waypoint_index]
        velocity_setpoint = calculate_velocity_setpoint(target_pos, pos)
        vx_setpoint, vy_setpoint, vz_setpoint = velocity_setpoint
    elif current_mode == FlightMode.RETURN_HOME:
        velocity_setpoint = calculate_velocity_setpoint(HOME_POSITION, pos)
        vx_setpoint, vy_setpoint, vz_setpoint = velocity_setpoint
    elif current_mode == FlightMode.LANDING:
        vx_setpoint = 0
        vy_setpoint = 0
        vz_setpoint = -min(LANDING_SPEED, (pos[2] - pid_z.setpoint) / 2)
    
    # 设置速度PID控制器的设定点
    pid_vx.set_setpoint(vx_setpoint)
    pid_vy.set_setpoint(vy_setpoint)
    pid_vz.set_setpoint(vz_setpoint)
    
    # 速度PID控制器输出
    fx = pid_vx.update(vel[0], DT)
    fy = pid_vy.update(vel[1], DT)
    fz = pid_vz.update(vel[2], DT) + MASS * GRAVITY

    # 计算姿态控制力矩
    tau_roll = pid_roll.update(att[0], DT)
    tau_pitch = pid_pitch.update(att[1], DT)
    tau_yaw = pid_yaw.update(att[2], DT)

    # 合力
    force = np.array([fx, fy, fz])
    torque = np.array([tau_roll, tau_pitch, tau_yaw])

    # 加速度
    acc = force / MASS - np.array([0, 0, GRAVITY])
    ang_acc = torque / INERTIA

    # 更新速度和位置
    vel += acc * DT
    pos += vel * DT
    
    # 防止降到地面以下
    if pos[2] < 0:
        pos[2] = 0
        vel[2] = 0

    # 更新角速度和姿态
    ang_vel += ang_acc * DT
    att += ang_vel * DT

    # 计算基于运动的姿态信息
    calculated_att = calculate_attitude_from_motion(vel, acc, DT)
    
    # 更新旋翼角度（模拟旋转）
    rotor_speed = 20 * DT  # 旋翼旋转速度
    rotor_angle += rotor_speed

    # 记录历史数据
    positions.append(pos.copy())
    attitudes.append(calculated_att.copy())
    flight_modes.append(current_mode)
    velocities.append(vel.copy())

    # 更新四旋翼模型
    body_faces, arms, rotors, landing_gear = create_quadcopter_model(pos, att, rotor_angle)
    
    # 更新机身
    body_collection.set_verts(body_faces)

    # 更新旋翼臂
    for i, arm_faces in enumerate(arms):
        arm_collections[i].set_verts(arm_faces)

    # 更新旋翼
    for i, rotor_faces in enumerate(rotors):
        rotor_collections[i].set_verts(rotor_faces)

    # 更新起落架
    for i, gear_faces in enumerate(landing_gear):
        landing_gear_collections[i].set_verts(gear_faces)

    # 更新姿态指示器
    axes = create_attitude_indicators(pos, att)
    x_axis.set_data([pos[0], axes[0,0]], [pos[1], axes[0,1]])
    x_axis.set_3d_properties([pos[2], axes[0,2]])
    y_axis.set_data([pos[0], axes[1,0]], [pos[1], axes[1,1]])
    y_axis.set_3d_properties([pos[2], axes[1,2]])
    z_axis.set_data([pos[0], axes[2,0]], [pos[1], axes[2,1]])
    z_axis.set_3d_properties([pos[2], axes[2,2]])

    # 更新轨迹线
    positions_array = np.array(positions)
    trajectory_line.set_data(positions_array[:, 0], positions_array[:, 1])
    trajectory_line.set_3d_properties(positions_array[:, 2])
    
    # 更新速度向量
    velocity_magnitude = np.linalg.norm(vel)
    velocity_vector.set_data([pos[0], pos[0]+vel[0]*0.2], [pos[1], pos[1]+vel[1]*0.2])
    velocity_vector.set_3d_properties([pos[2], pos[2]+vel[2]*0.2])

    # 更新模式文本
    mode_text.set_text(f'Test Mode: {test_mode.name}\nFlight Mode: {current_mode.name}')
    speed_text.set_text(f'Speed: {velocity_magnitude:.2f} m/s')

    return (body_collection, *arm_collections, *rotor_collections, *landing_gear_collections, 
            x_axis, y_axis, z_axis, trajectory_line, velocity_vector, mode_text, speed_text)

def save_data_to_file(positions, attitudes, filename="flight_data.txt"):
    with open(filename, 'w') as f:
        f.write("Time(s), Pos_X, Pos_Y, Pos_Z, Att_Roll, Att_Pitch, Att_Yaw\n")
        for i, (pos, att) in enumerate(zip(positions, attitudes)):
            time = i * DT  # 假设DT是时间步长
            f.write(f"{time:.4f}, {pos[0]:.4f}, {pos[1]:.4f}, {pos[2]:.4f}, "
                   f"{att[0]:.4f}, {att[1]:.4f}, {att[2]:.4f}\n")
    print(f"Flight data saved to {filename}")

# 创建动画
ani = FuncAnimation(fig, update, frames=int(SIM_TIME/DT), interval=DT*1000, blit=True)
plt.legend()
plt.tight_layout()
plt.show()
# 动画结束后保存数据
save_data_to_file(positions, attitudes)














#未用代码
def create_quadcopter_vertices(pos, att):
    """创建四旋翼的顶点"""
    # 创建基础顶点（未旋转）
    vertices = np.array([
        [ARM_LENGTH, 0, 0],    # 右前
        [0, ARM_LENGTH, 0],    # 前
        [-ARM_LENGTH, 0, 0],   # 左前
        [0, -ARM_LENGTH, 0],   # 后
    ])
    
    # 旋转矩阵
    roll, pitch, yaw = att
    Rx = np.array([[1, 0, 0],
                   [0, np.cos(roll), -np.sin(roll)],
                   [0, np.sin(roll), np.cos(roll)]])
    
    Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)],
                   [0, 1, 0],
                   [-np.sin(pitch), 0, np.cos(pitch)]])
    
    Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0],
                   [np.sin(yaw), np.cos(yaw), 0],
                   [0, 0, 1]])
    
    R = Rz @ Ry @ Rx
    
    # 应用旋转
    vertices = vertices @ R.T
    
    # 添加位置偏移
    vertices += pos
    
    return vertices